from eden.evo.cell import NonsexualCell, SexualCell
from eden.evo.chromosome import Chromosome
from eden.evo.context import Context
from eden.evo.mutate import ChangeMutator, InsertMutator, DeleteMutator, \
    OriginateMutator
from eden.evo.reproduction import NonsexualReproducer, SexualReproducer
from eden.evo.transcode import Transcoder
from eden.evo.translate import Translator
import doctest
import os

class Pool:
    MUTATE_RATIO = 0.01

    def __init__(self, cell_class, reproducer_class,
            cell_count, init_chromosome_len, encoding,
            evaluator, target_hitratio = 0.9):
        self.__init_chromosome_len = init_chromosome_len
        self.__transcoder = Translator(encoding)
        self.__evaluator = evaluator
        self.__target_hitratio = target_hitratio
        self.generation = 0
        self.fitness = 0.0
        self.hitratio = 0.0
        self.cells = [cell_class.create_null_cell(init_chromosome_len, self.__transcoder) for i in range(cell_count)] #@UnusedVariable
        self.reproducer = self._create_reproducer(reproducer_class, encoding)
        self.events = []
    
    def _create_reproducer(self, reproducer_class, encoding):
        mutators = (ChangeMutator(Pool.MUTATE_RATIO, encoding),
                    InsertMutator(Pool.MUTATE_RATIO, encoding),
                    DeleteMutator(Pool.MUTATE_RATIO, encoding),
                    OriginateMutator(1 - self.__target_hitratio, encoding, Chromosome.create_null_chromosome(self.__init_chromosome_len)))
        return reproducer_class(mutators)
    
    def _runcell(self, cell, input, expected):
        '''
        >>> from eden.evo.aminos import *
        >>> from eden.evo.evaluator import *
        >>> encoding = { 0 : NullAmino }
        >>> pool = Pool(NonsexualCell, NonsexualReproducer, 0, 0, encoding, SimpleEvaluator())
        >>> cell = NonsexualCell.create_null_cell(1, Translator(encoding))
        >>> pool._runcell(cell, (1, ), (2, ))
        (0.5, False)
        >>> pool._runcell(cell, (1, ), (1, ))
        (1.0, True)
        '''
        context = Context.from_data(input)
        cell.run(context)
        output = context.getregisters()[:len(expected)]
        fitness = 0.0
        hitted = False
        if expected is not None:
            fitness = self.__evaluator.evaluate(expected, output)
            if output == expected:
                hitted = True
        return fitness, hitted
        
    def run(self, data):
        '''
        >>> from eden.evo.aminos import *
        >>> from eden.evo.evaluator import *
        >>> encoding = { 0 : NullAmino }
        >>> pool = Pool(NonsexualCell, NonsexualReproducer, 2, 1, encoding, SimpleEvaluator())
        >>> data = (([1], [2]), ([1], [2]), ([1], [2]))
        >>> pool.run(data)
        (0.5, 0.0)
        >>> [cell.fitness for cell in pool.cells]
        [0.5, 0.5]
        >>> [cell.hitratio for cell in pool.cells]
        [0.0, 0.0]
        '''
        for cell in self.cells:
            fitness_sum = 0.0
            hits = 0
            for (input, expected) in data:
                fitness, hitted = self._runcell(cell, input, expected)
                fitness_sum += fitness
                if hitted:
                    hits += 1
            cell.fitness = fitness_sum / len(data)
            cell.hitratio = float(hits) / len(data)
        fitness_avg = sum([cell.fitness for cell in self.cells]) / len(self.cells)
        hitness_avg = sum([cell.hitratio for cell in self.cells]) / len(self.cells)
        return fitness_avg, hitness_avg
        
    def reproduce(self):
        self.generation += 1
        self.reproducer.reproduce(self.cells, self.generation)
        
    def save_to_file(self, dir):
        if not os.path.exists(dir):
            os.mkdir(dir)
        Transcoder().save_to_file(self.cells, dir, self.generation)
        
    def evolve(self, data, max_generations = -1):
        while True:
            self.fitness, self.hitratio = self.run(data)
            self.__fire_events()
            if self.hitratio > self.__target_hitratio:
                return
            if self.generation == max_generations:
                return
            self.reproduce()
            
    def add_event(self, event):
        self.events.append(event)
        
    def remove_event(self, event):
        self.events.remove(event)
        
    def __fire_events(self):
        for event in self.events:
            event(self)
        
class NonsexualPool(Pool):
    def __init__(self, cell_count, init_chromosome_len,
            encoding, evaluator):
        Pool.__init__(self, NonsexualCell,
                NonsexualReproducer, cell_count,
                init_chromosome_len, encoding, evaluator)
        
class SexualPool(Pool):
    def __init__(self, cell_count, init_chromosome_len, encoding, evaluator, save_dir):
        Pool.__init__(self, SexualCell, SexualReproducer,
                cell_count, init_chromosome_len, encoding,
                evaluator, save_dir)
        
if __name__ == '__main__':
    doctest.testmod()
